# -*- coding: utf-8 -*-
import rospy
from mavros_msgs.msg import GlobalPositionTarget, State, PositionTarget
from mavros_msgs.srv import CommandBool, CommandTOL, SetMode 
from geometry_msgs.msg import PoseStamped, Twist
from sensor_msgs.msg import Imu, NavSatFix
from std_msgs.msg import Float32, Float64, String
import time
import math
import thread
from pyquaternion import Quaternion


class Controller:
    def __init__(self):
        #定义成员变量
        self.local_position = None
        self.local_enu_position = None

        self.state = None
        self.global_position = None
        self.current_heading = None

        self.current_target_position = None
        self.current_target_pos = None

        self.received_imu = False
        self.arm_state = False
        self.offb_state = False

        self.takeoff_height = 1.8


        #订阅
        self.local_position_subscrib = rospy.Subscriber("mavros/local_position/pose", PoseStamped, self.callback_local_position_subscribe)
        self.state_subscrib = rospy.Subscriber("mavros/state", State, self.callback_state_subscrib)
        self.global_position_subscrib = rospy.Subscriber("mavros/global_position/global", NavSatFix, self.callback_global_position_subscrib)
        self.imu_subscrib = rospy.Subscriber("mavros/imu/data", Imu, self.callback_imu_subscrib)

        #发布
        self.local_target_publish = rospy.Publisher("mavros/setpoint_raw/local", PositionTarget, queue_size=1)
        self.local_position_publish = rospy.Publisher("mavros/setpoint_position/local", PoseStamped, queue_size=1)

        #服务
        self.arm_service = rospy.ServiceProxy("mavros/cmd/arming", CommandBool)
        self.flight_mode_service = rospy.ServiceProxy("mavros/set_mode", SetMode)

        print("px4 Controller initialized!")
    
    #回调函数,python的函数参数不用定义数据类型
    def callback_local_position_subscribe(self, msg):
        self.local_position = msg
        self.local_enu_position = msg
    
    def callback_state_subscrib(self, msg):
        self.state = msg.mode
    
    def callback_global_position_subscrib(self, msg):
        self.global_position = msg
    
    def callback_imu_subscrib(self, msg):
        global global_imu, current_heading
        self.imu = msg
        self.current_heading = self.q2yaw(self.imu.orientation)
        
        self.received_imu = True
    
    def construct_target(self, x, y, z, yaw, yaw_rate=1):
        target_raw_position = PositionTarget()
        target_raw_position.header.stamp = rospy.Time.now()
    #9是FRAME_BODY_OFFSET_NED = 9
        target_raw_position.coordinate_frame = 1

        target_raw_position.position.x = x
        target_raw_position.position.y = y
        target_raw_position.position.z = z

        target_raw_position.type_mask = PositionTarget.IGNORE_VX + PositionTarget.IGNORE_VY + PositionTarget.IGNORE_VZ + PositionTarget.IGNORE_AFX + PositionTarget.IGNORE_AFY + PositionTarget.IGNORE_AFZ + PositionTarget.FORCE

        target_raw_position.yaw = yaw
        target_raw_position.yaw_rate = yaw_rate
        return target_raw_position
    
    def q2yaw(self, q):
        if isinstance(q, Quaternion):
            rotate_z_rad = q.yaw_pitch_roll[0]
        else:
            q_ = Quaternion(q.w, q.x, q.y, q.z)
            rotate_z_rad = q_.yaw_pitch_roll[0]
        
        return rotate_z_rad
    
    #发送解锁命令，解锁成功服务端返回true，失败返回false
    def arm(self):
        if self.arm_service(True):
            return True
        else:
            print("arming false")
            return False
    
    def disarm(self):
        if self.arm_service(False):
            return True
        else:
            print("disarm failed")
    
    #发送切换到OFFBOARD模式的指令，切换成功返回True，失败返回False
    def offboard(self):
        if self.flight_mode_service(custom_mode = "OFFBOARD"):
            return True
        else:
            print("offboard false")
            return False

    #检查飞机是否起飞，是否有离地高度
    def takeoff_detection(self):
        if self.local_position.pose.position.z > 0.1:
            return True
        else:
            return False
    
    def is_arrived(self, x_, y_):
        #检查是否到达指定点
        if ((abs(self.local_position.pose.position.y - y_) < 0.3) and (abs(self.local_position.pose.position.x - x_) < 0.3)):
            print("arrived")
            print(self.local_position)
            return True
        else:
            print("not arriveed")
            print(self.local_position)
            return False
###############################################################################
#使用地球坐标下的位置点来发送目标点
    def change_to_next_point(self, x_, y_, z_, yaw=0, yaw_rate=0):
        next_point = PoseStamped()
        next_point.header.stamp = rospy.Time.now()

        next_point.pose.position.x = x_
        next_point.pose.position.y = y_
        next_point.pose.position.z = z_

        #next_point.type_mask = PositionTarget.IGNORE_VX + PositionTarget.IGNORE_VY + PositionTarget.IGNORE_VZ + PositionTarget.IGNORE_AFX + PositionTarget.IGNORE_AFY + PositionTarget.IGNORE_AFZ + PositionTarget.FORCE
        
        #next_point.yaw = yaw
        #next_point.yaw_rate = yaw_rate

        return next_point
    
###############################################################################

    def start(self):
        rospy.init_node("offb_node_python")
        #解锁前要给发布目标
        self.current_target_position = self.construct_target(0, 0, self.takeoff_height, self.current_heading)
        for i in range(10):
            self.local_target_publish.publish(self.current_target_position)
            #self.local_position_publish.publish()
            self.arm_state = self.arm()
            self.offb_state = self.offboard()
        
        if self.takeoff_detection():
            print("Takeoff true")
        else:
            print("takeoff false")
        
        #解锁，切换到offboard以及程序没有退出，开始切入外部控制的发布状态
        height = 10
        x=[0.0, 15.0, 15.0, 0.0]
        y=[0.0, 0.0, 15.0, 15.0]
        i=0

        self.current_target_pos = self.change_to_next_point(0 , 0, height)
        self.local_position_publish.publish(self.current_target_pos)

        while self.arm_state and self.offb_state and (rospy.is_shutdown() is False):
            if self.is_arrived(x[i],y[i]):
                #self.current_target_position = self.construct_target(self.current_target_position.position.x,self.current_target_position.position.y+10,self.current_target_position.position.z,0, 1)  
                if i < 3:
                    i=i+1
                else:
                    i=0
                self.current_target_pos = self.change_to_next_point(x[i],y[i],height)
                print("arrive:")
                print(i)

            self.local_position_publish.publish(self.current_target_pos)

            if (self.state is "LAND") and (self.local_position.pose.position.z < 0.15):
                if(self.disarm()):
                    self.state = "DISARMED"
            
            time.sleep(0.1)

if __name__ == "__main__":
    con = Controller()
    con.start()


